import lightgbm as lgb
from sklearn.metrics import accuracy_score
from dataPretreat import data_pretreat

if __name__ == '__main__':
    # 加载数据

    iszt_train = data_pretreat(path='../resource/KDDTrain+.csv')
    iszt_test = data_pretreat(path='../resource/KDDTest+.csv')
    # 划分训练集和测试集

    X_train= iszt_train[['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
                        'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
                        'dst_host_rerror_rate']]
    y_train= iszt_train['label']
    X_test = iszt_test[
        ['duration', 'protocol_type', 'service', 'src_bytes', 'dst_bytes', 'wrong_fragment', 'serror_rate',
        'dst_host_srv_count', 'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate',
        'dst_host_rerror_rate']]
    y_test = iszt_test['label']

    # 转换为Dataset数据格式

    train_data = lgb.Dataset(X_train, label=y_train)
    validation_data = lgb.Dataset(X_test, label=y_test)

    # 参数

    params = {
        'learning_rate': 0.1,
        'lambda_l1': 0.1,
        'lambda_l2': 0.2,
        'max_depth': 4,
        'objective': 'multiclass',  # 目标函数
        'num_class': 37,
    }

    # 模型训练
    gbm = lgb.train(params, train_data, valid_sets=[validation_data])

    # 模型预测
    y_pred = gbm.predict(X_test)
    y_pred = [list(x).index(max(x)) for x in y_pred]
    print(y_pred)

    # 模型评估
    acc=accuracy_score(y_test, y_pred)*100
    print("准确率百分之%.3f" % acc)